"""
PyG兼容层中的Graph类实现

这个模块提供了兼容DGL的Graph类实现，使用PyG作为后端
"""

import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu
from torch_geometric.data import Data

# 通过条件导入避免循环引用
try:
    import dgl
except ImportError:
    dgl = None


class Graph:
    """
    PyG兼容的图结构，API设计与DGL兼容
    """
    def __init__(self, edges=None, num_nodes=None, batch_size=1):
        """
        创建图结构
        
        参数:
            edges: 元组 (src, dst) 表示边的源节点和目标节点，如果为None则创建空图
            num_nodes: 节点数量，如果为None则自动推断
            batch_size: 图的批处理大小，默认为1
        """
        if edges is None:
            # 创建空图（兼容PyG的Batch系统）
            self.edge_index = torch.empty((2, 0), dtype=torch.long)
            num_nodes = 0 if num_nodes is None else num_nodes
        else:
            src, dst = edges
            self.edge_index = torch.stack([src, dst])
            
            if num_nodes is None:
                num_nodes = max(self.edge_index.max().item() + 1, 
                                 (self.edge_index[0].max() + 1).item(),
                                 (self.edge_index[1].max() + 1).item()) if self.edge_index.numel() > 0 else 0
        
        self._num_nodes = num_nodes
        self._edge_data = {}
        self._node_data = {}
        self._device = torch.device(self.edge_index.device)
        self._batch_size = batch_size
        self._num_edges = self.edge_index.size(1)
        
        # 添加模拟DGL内部结构的属性
        class _DGLGraph:
            def __init__(self, parent_graph):
                self._parent = parent_graph
                self.edge_index = parent_graph.edge_index
                self._num_nodes = parent_graph._num_nodes
                self._etypes = [0]  # 默认边类型
                self._ntypes = [0]  # 默认节点类型
                self._sparse_matrices = {}  # 缓存稀疏矩阵
                self._batch_num_nodes = [parent_graph._num_nodes]  # 每个批次的节点数
                self._batch_num_edges = [parent_graph.edge_index.size(1)]  # 每个批次的边数
                
            def number_of_etypes(self):
                return 1
                
            def number_of_edges(self, *args, **kwargs):
                # 忽略所有参数，模拟DGL的行为
                return self.edge_index.size(1)
                
            def number_of_nodes(self, *args, **kwargs):
                # 忽略所有参数，模拟DGL的行为
                return self._num_nodes
                
            def __repr__(self):
                return f"_DGLGraph(num_nodes={self._num_nodes}, num_edges={self.edge_index.size(1)})"
                
            # 增加与DGL兼容的方法
            def etypes(self):
                return self._etypes
                
            def ntypes(self):
                return self._ntypes
                
            def canonical_etypes(self):
                # DGL中的标准边类型格式 (src_type, edge_type, dst_type)
                return [(0, 0, 0)]
                
            def to_simple(self, *args, **kwargs):
                # 返回简化图（自身，因为我们已经是简单图）
                return self
                
            def adjacency_matrix(self, *args, **kwargs):
                # 返回稀疏邻接矩阵（简单实现）
                indices = self.edge_index
                values = torch.ones(indices.size(1), device=indices.device)
                return torch.sparse.FloatTensor(
                    indices, values, 
                    torch.Size([self._num_nodes, self._num_nodes])
                )
                
            def incidence_matrix(self, *args, **kwargs):
                # 返回关联矩阵（简单实现）
                src, dst = self.edge_index
                indices = torch.stack([
                    torch.arange(src.size(0), device=src.device),
                    src
                ])
                values = torch.ones(indices.size(1), device=indices.device)
                size = torch.Size([indices.size(1), self._num_nodes])
                return torch.sparse.FloatTensor(indices, values, size)
                
            def in_degrees(self, v=None):
                # 计算入度
                dst = self.edge_index[1]
                if v is not None:
                    mask = dst == v
                    return torch.sum(mask)
                else:
                    return torch.bincount(dst, minlength=self._num_nodes)
                
            def out_degrees(self, v=None):
                # 计算出度
                src = self.edge_index[0]
                if v is not None:
                    mask = src == v
                    return torch.sum(mask)
                else:
                    return torch.bincount(src, minlength=self._num_nodes)
            
            # 添加图消息传递相关的方法
            def send_and_recv(self, src_nodes, dst_nodes, message_func, reduce_func, apply_node_func=None):
                """模拟DGL的send_and_recv操作"""
                # 在SE3Transformer中可能会被调用
                if isinstance(src_nodes, int) or isinstance(dst_nodes, int):
                    # 转为列表
                    if isinstance(src_nodes, int):
                        src_nodes = [src_nodes]
                    if isinstance(dst_nodes, int):
                        dst_nodes = [dst_nodes]
                
                # 获取源节点和目标节点的特征
                src_feats = {k: v[src_nodes] for k, v in self._parent._node_data.items() if k in message_func.keys()}
                
                # 应用消息函数
                msgs = message_func(src_feats)
                
                # 应用归约函数
                reduced = reduce_func(msgs)
                
                # 更新目标节点特征
                if apply_node_func:
                    dst_feats = {k: v[dst_nodes] for k, v in self._parent._node_data.items() if k in apply_node_func.keys()}
                    dst_feats.update(reduced)
                    apply_node_func(dst_feats)
                
                return reduced
                
            # 添加数据访问方法
            def nodes(self, *args, **kwargs):
                return torch.arange(self._num_nodes, device=self._parent.device)
                
            def edges(self, *args, **kwargs):
                src, dst = self.edge_index
                if args:
                    # 如果指定了src或dst，返回特定边
                    u, v = args
                    if u is not None:
                        mask_src = src == u
                    else:
                        mask_src = torch.ones_like(src, dtype=torch.bool)
                        
                    if v is not None:
                        mask_dst = dst == v
                    else:
                        mask_dst = torch.ones_like(dst, dtype=torch.bool)
                        
                    mask = mask_src & mask_dst
                    return src[mask], dst[mask]
                return src, dst
                
            def ndata(self, *args, **kwargs):
                # 在SE3Transformer中可能会被调用
                class NDataView:
                    def __init__(self, graph_ref):
                        self.graph_ref = graph_ref
                    
                    def __getitem__(self, key):
                        # 获取节点特征
                        return self.graph_ref._parent._node_data.get(key, None)
                    
                    def __setitem__(self, key, value):
                        # 设置节点特征
                        self.graph_ref._parent._node_data[key] = value
                
                return NDataView(self)
                
            def edata(self, *args, **kwargs):
                # 在SE3Transformer中可能会被调用
                class EDataView:
                    def __init__(self, graph_ref):
                        self.graph_ref = graph_ref
                    
                    def __getitem__(self, key):
                        # 获取边特征
                        return self.graph_ref._parent._edge_data.get(key, None)
                    
                    def __setitem__(self, key, value):
                        # 设置边特征
                        self.graph_ref._parent._edge_data[key] = value
                
                return EDataView(self)
                
            def __class_getitem__(self, *args, **kwargs):
                # 支持DGL中的图类型标注
                return self.__class__
                
            def __call__(self, *args, **kwargs):
                # 特定类型的图调用
                return self
                
            def update_all(self, message_func, reduce_func, apply_node_func=None):
                # 模拟DGL的update_all操作
                src, dst = self.edge_index
                
                # 获取源节点特征
                src_data = {k: v[src] for k, v in self._parent._node_data.items()}
                
                # 获取边特征
                edge_data = self._parent._edge_data
                
                # 合并特征
                msg_data = {}
                msg_data.update(src_data)
                msg_data.update(edge_data)
                
                # 应用消息函数
                msgs = message_func(msg_data)
                
                # 按目标节点聚合
                reduced = {}
                for k, v in msgs.items():
                    reduced[k] = torch.zeros((self._num_nodes, v.size(-1)), device=v.device)
                    reduced[k].index_add_(0, dst, v)
                
                # 应用节点函数
                if apply_node_func:
                    node_data = {k: v for k, v in self._parent._node_data.items()}
                    node_data.update(reduced)
                    apply_node_func(node_data)
            
            def is_homogeneous(self):
                # 简单图一般是同构的
                return True
            
            def is_multigraph(self):
                # 我们当前不支持多重图
                return False
                
            def has_node(self, vid):
                # 检查节点是否存在
                return 0 <= vid < self._num_nodes
                
            def has_edge_between(self, u, v):
                # 检查边是否存在
                src, dst = self.edge_index
                for i in range(src.size(0)):
                    if src[i] == u and dst[i] == v:
                        return True
                return False
                
            # 支持获取和修改特征
            def batch_num_nodes(self):
                # 返回每个批次的节点数量
                return self._batch_num_nodes
                
            def batch_num_edges(self):
                # 返回每个批次的边数量
                return self._batch_num_edges
        
        # 创建内部DGL图对象并设置循环引用
        self._graph = _DGLGraph(self)
    
    def to(self, device):
        """将图移动到指定设备"""
        # 确保device是torch.device对象
        device = torch.device(device)
        
        self.edge_index = self.edge_index.to(device)
        
        # 递归地将边特征移动到设备
        for k in self._edge_data:
            if isinstance(self._edge_data[k], torch.Tensor):
                self._edge_data[k] = self._edge_data[k].to(device)
            elif isinstance(self._edge_data[k], dict):
                self._edge_data[k] = {kk: vv.to(device) if isinstance(vv, torch.Tensor) else vv 
                                     for kk, vv in self._edge_data[k].items()}
            elif isinstance(self._edge_data[k], list):
                self._edge_data[k] = [vv.to(device) if isinstance(vv, torch.Tensor) else vv 
                                     for vv in self._edge_data[k]]
        
        # 递归地将节点特征移动到设备
        for k in self._node_data:
            if isinstance(self._node_data[k], torch.Tensor):
                self._node_data[k] = self._node_data[k].to(device)
            elif isinstance(self._node_data[k], dict):
                self._node_data[k] = {kk: vv.to(device) if isinstance(vv, torch.Tensor) else vv 
                                     for kk, vv in self._node_data[k].items()}
            elif isinstance(self._node_data[k], list):
                self._node_data[k] = [vv.to(device) if isinstance(vv, torch.Tensor) else vv 
                                     for vv in self._node_data[k]]
        
        self._device = device
        return self
    
    def edges(self):
        """
        返回图的源节点和目标节点
        
        返回:
            (src, dst): 源节点和目标节点的元组
        """
        return self.edge_index[0], self.edge_index[1]
    
    @property
    def device(self):
        """获取图所在的设备"""
        return self._device
    
    @property
    def edata(self):
        """获取边特征字典，兼容DGL的接口风格"""
        return EdgeDataView(self)
    
    @property
    def ndata(self):
        """获取节点特征字典，兼容DGL的接口风格"""
        return NodeDataView(self)
    
    @property
    def batch_size(self):
        """获取图的批处理大小"""
        return self._batch_size
    
    @batch_size.setter
    def batch_size(self, value):
        """设置图的批处理大小"""
        self._batch_size = value
    
    def num_edges(self):
        """获取边的数量"""
        return self._num_edges
    
    def num_nodes(self):
        """获取节点的数量"""
        return self._num_nodes
    
    def to_pyg_data(self):
        """转换为PyG的Data对象"""
        data = Data(
            edge_index=self.edge_index,
            num_nodes=self._num_nodes
        )
        
        # 添加边特征
        for k, v in self._edge_data.items():
            data[f'edge_{k}'] = v
            
        # 添加节点特征
        for k, v in self._node_data.items():
            data[k] = v
            
        return data
    
    def is_dgl_graph(self):
        """检查是否是DGL图"""
        return False


class EdgeDataView:
    """边特征视图，模拟DGL的edata接口"""
    def __init__(self, graph):
        self.graph = graph
    
    def __getitem__(self, key):
        return self.graph._edge_data.get(key, None)
    
    def __setitem__(self, key, value):
        self.graph._edge_data[key] = value
    
    def items(self):
        return self.graph._edge_data.items()


class NodeDataView:
    """节点特征视图，模拟DGL的ndata接口"""
    def __init__(self, graph):
        self.graph = graph
    
    def __getitem__(self, key):
        return self.graph._node_data.get(key, None)
    
    def __setitem__(self, key, value):
        self.graph._node_data[key] = value
    
    def items(self):
        return self.graph._node_data.items()


# 兼容DGL的图创建函数
def graph(edges=None, num_nodes=None, batch_size=1):
    """
    创建图对象，兼容DGL的dgl.graph接口
    
    参数:
        edges: 元组 (src, dst) 表示边的源节点和目标节点，如果为None则创建空图
        num_nodes: 节点数量，如果为None则自动推断
        batch_size: 图的批处理大小，默认为1
    
    返回:
        Graph对象
    """
    return Graph(edges, num_nodes, batch_size)


# 判断是否为DGL图
def is_dgl_graph(graph):
    """
    判断一个图是否为DGL图，可通过环境变量控制
    
    参数:
        graph: 图对象
    
    返回:
        True如果是DGL图或者环境变量要求使用DGL，否则False
    """
    import os
    
    # 检查环境变量
    force_dgl = os.environ.get('USE_DGL', '').lower() in ('true', '1', 'yes')
    force_pyg = os.environ.get('FORCE_PYG', '').lower() in ('true', '1', 'yes')
    
    if force_dgl and not force_pyg:
        # 如果强制使用DGL，对非DGL图可以尝试转换
        if dgl is None:
            return False
        if isinstance(graph, dgl.DGLGraph):
            return True
        # 对于其他类型的图可以添加尝试转换的逻辑
        try:
            # 这里只是检查类型，不进行实际的转换
            print(f"警告: 环境变量要求使用DGL，但输入不是DGL图，尝试适配...")
            return False  # 在这里返回False，让调用方处理转换逻辑
        except Exception:
            return False
    elif force_pyg:
        # 如果强制使用PyG，则不使用DGL
        return False
    else:
        # 默认行为：根据实际类型判断
        if dgl is None:
            return False
        return isinstance(graph, dgl.DGLGraph) 